package com.gitee.xmhzzz.component.rocketmq.support;


import cn.hutool.core.collection.CollUtil;
import com.gitee.xmhzzz.component.rocketmq.context.TraceContext;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.rocketmq.client.hook.ConsumeMessageContext;
import org.apache.rocketmq.client.hook.ConsumeMessageHook;
import org.apache.rocketmq.client.hook.SendMessageContext;
import org.apache.rocketmq.client.hook.SendMessageHook;
import org.apache.rocketmq.client.producer.DefaultMQProducer;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.spring.support.DefaultRocketMQListenerContainer;
import org.slf4j.MDC;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.event.EventListener;

import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * @ClassName pu
 * @Description
 * @Author wzq
 * @Date 2023/12/4 15:46
 * @Version 1.0
 */
public class RocketMqTraceListener {

    @EventListener
    public void onApplicationEvent(ApplicationReadyEvent event) {
        ConfigurableApplicationContext applicationContext = event.getApplicationContext();

        Map<String, DefaultMQProducer> defaultMQProducerMap = applicationContext.getBeansOfType(DefaultMQProducer.class);
        if(CollUtil.isNotEmpty(defaultMQProducerMap)) {
            for (Map.Entry<String, DefaultMQProducer> map : defaultMQProducerMap.entrySet()) {
                DefaultMQProducer value = map.getValue();
                value.getDefaultMQProducerImpl().registerSendMessageHook(new SendMessageHook() {
                    @Override
                    public String hookName() {
                        return "traceSendMessageHookImpl";
                    }

                    @Override
                    public void sendMessageBefore(SendMessageContext sendMessageContext) {
                        String traceId = (String) MDC.get("X-B3-TraceId");
                        traceId = StringUtils.isNotBlank(traceId) ? traceId : TraceContext.getTraceId();
                        String spanId = (String) MDC.get("X-B3-SpanId");
                        spanId = StringUtils.isNotBlank(spanId) ? spanId : TraceContext.getSpanId();
                        Message message = sendMessageContext.getMessage();
                        message.putUserProperty("traceId", traceId);
                        message.putUserProperty("spanId", spanId);
                    }

                    @Override
                    public void sendMessageAfter(SendMessageContext sendMessageContext) {

                    }
                });
            }
        }

        Map<String, DefaultRocketMQListenerContainer> beansOfType = applicationContext.getBeansOfType(DefaultRocketMQListenerContainer.class);
        if(beansOfType.isEmpty()){
            return;
        }
        for (Map.Entry<String, DefaultRocketMQListenerContainer> stringDefaultRocketMQListenerContainerEntry : beansOfType.entrySet()) {
            DefaultRocketMQListenerContainer value = stringDefaultRocketMQListenerContainerEntry.getValue();
            value.getConsumer().getDefaultMQPushConsumerImpl().registerConsumeMessageHook(new ConsumeMessageHook() {
                @Override
                public String hookName() {
                    return null;
                }
                @Override
                public void consumeMessageBefore(ConsumeMessageContext context) {
                    String traceId = null;
                    String spanId = null;

                    List<MessageExt> msgList = context.getMsgList();
                    if(CollectionUtils.isNotEmpty(msgList) && msgList.size() == 1){
                        MessageExt messageExt = msgList.get(0);
                        traceId = messageExt.getProperty("traceId");
                        spanId = messageExt.getProperty("spanId");
                    }
                    // 可以取生产者传递
                    if (StringUtils.isEmpty(traceId)) {
                        // 可以取生产者传递
                        Map<String, String> props = context.getProps();
                        traceId = props.get("traceId");
                        spanId = props.get("spanId");
                    }
                    // 自定义
                    if (StringUtils.isEmpty(traceId)) {
                        traceId = UUID.randomUUID().toString().replaceAll("-", "");
                    }
                    // 自定义
                    if (StringUtils.isEmpty(spanId)) {
                        spanId = UUID.randomUUID().toString().replaceAll("-", "");
                    }
                    MDC.put("X-B3-TraceId", traceId);
                    MDC.put("X-B3-SpanId", spanId);

                }

                @Override
                public void consumeMessageAfter(ConsumeMessageContext context) {
                    MDC.clear();
                }
            });
        }
    }
}
